﻿#-------------------------------------------------------------------------------
# Name:        actions
# Purpose:
#
# Author:      SUPER_紫电  QQ,911344755
#
# Created:     28/01/2023
# Copyright:  (c) SUPER_紫电 2022 All rights reserved.
# Licence:    <Modified BSD License>
#-------------------------------------------------------------------------------

import datetime, time
import colorconsole, common, mutex, taskfile
from taskfile import *
from vsc import *

cprint  = colorconsole.cprint
errinfo = colorconsole.errinfo

def get_ident():
    print(get_identify())

def get_vscstat():
    print(dec_vsc_status())

def beep():
    freq = 1193180000 / 440000 # 440Hz
    outp(0x43, 0xB6)
    outp(0x42, int(freq) & 0xFF)
    outp(0x42, int(freq) >> 9)
    sleep(0.1)
    outp(0x61, inp(0x61) | 0x03)
    sleep(1)
    outp(0x61, inp(0x61) & 0xFC)

def poll_vscstat():
    if not get_vsc_status():
        print('Error executing command')
        return
    drv_state1 = {
        0: 'Idle',
        1: 'Standby',
        2: 'Sleep',
        3: 'Background DST',
        4: 'Background Data Lifeguard',
        5: 'Background VSC Command',
        6: 'PST test in progress',
    }
    vscstatus = vsctype.vsc_status
    cur_state = vscstatus.CurState
    print('Current Drive State 1 = %d %s' % (cur_state, drv_state1[cur_state]))
    if cur_state != 6:
        return
    start = time()
    while cur_state == 6:
        cls()
        runtime = time()-start
        if runtime//3600>24:
            runtime = '%d-%02d:%02d:%02d' % (runtime//3600//24,runtime//3600%24,runtime%3600//60,runtime%60)
        else:
            runtime = '%02d:%02d:%02d' % (runtime//3600, runtime%3600//60,runtime%60)
        result  = 'PST test in progress - %s' % runtime
        #result  = 'PST test in progress - {:0>8}'.format(str(datetime.timedelta(seconds=time()-start)))
        result += '\n----------------------------------------------------------'
        result += '\nPTM File Id              = 0x%02X' % vscstatus.PTMId
        result += '\nPTM Test Id              = 0x%02X' % vscstatus.TestId
        result += '\nCurrent VSC Command      = 0x%02X' % vscstatus.ActionCode
        result += '\nExtended Error Status    = 0x%04X' % vscstatus.ExtErr
        result += '\nExtended Error Code      = 0x%04X' % vscstatus.ExtErrCode
        if vscstatus.PTMId == 0xB9 or vscstatus.PTMId == 0xBA:
            buf = vscstatus.CurLBA
            lba = 0
            for i in range(6):
                lba += buf[i]<<(i*8)
            result += '\nCurrent LBA              = %d' % lba
            prg = vscstatus.BGProgress
            result += '\nProgress Indicator       =  %d.%02d%%'%(prg//10,prg%10)
        result += '\nCurrent Virtual Cylinder = %d' % vscstatus.CurVirCyl
        result += '\nCurrent Virtual Head     = %d' % vscstatus.CurVirHd
        result += '\nCurrent Data Zone        = %d' % vscstatus.CurZone
        result += '\nDrive Temperature (avg)  = %d' % vscstatus.DTemp
        result += '\n----------------------------------------------------------'
        if not get_vsc_status(): break
        vscstatus = vsctype.vsc_status
        cur_state = vscstatus.CurState
        ptm_id = vscstatus.PTMId
        print(result)
        if ptm_id == 0xFF22 or ptm_id == 0xFFF1:
            break
        sleep(0.5)
    if vsctype.ext_err == 0:
        print('Completed Successfully')
    else:
        errinfo(dec_err_code())
    beep()

def vscid():
    common.workdir = ''
    drv_tbl = get_drv_tbl()
    if drv_tbl != '':
        print(drv_tbl)
    if common.workdir != '':
        path = common.get_startup_dir()
        common.workdir = path + '\\Work\\' + common.workdir
        print('\n[Work Dir]\n' + common.workdir)
        if not os.path.exists(common.workdir):
            os.makedirs(common.workdir)
        filename = common.workdir + '\\vscid.txt'
        if os.path.exists(filename):
            return
        with open(filename, 'w', encoding='utf-8') as file:
            file.write(drv_tbl)

def spin_down_reset():
    if not spindown_reset():
        get_err_code()
        errinfo('Reset Drive - ' + dec_err_code())
        return
    result = 'Reset Drive - please wait'
    print(result, end='\r')
    start = time()
    while reg_status(BSY):
        if time() - start > 30:
            print(result + 'time out!', end='\r\n')
            return
        result += '.'
        print(result, end='\r')
        sleep(1)
    print(result + 'done.', end='\r\n')

def servo_spin_down():
    sv_spindn()
    get_err_code()
    print('Servo Spin Down - ' + dec_err_code())

def servo_spin_up():
    if not sv_spinup():
        get_err_code()
        errinfo('Servo Spin Up - ' + dec_err_code())
        return
    result = 'Servo Spin Up - please wait'
    print(result, end='\r')
    start = time()
    while reg_status(BSY):
        if time() - start > 30:
            print(result + 'time out!', end='\r\n')
            return
        result += '.'
        print(result, end='\r')
        sleep(1)
    print(result + 'done.', end='\r\n')

def soft_reset():
    port = ata_port['base']
    dev = 0x50 if ata_port['dev'] == 0x40 else 0x40
    mutex_name = 'trex_mutex_0x%04x_0x%02x' % (port, dev)
    if mutex.exist_mutex(mutex_name):
        return
    taskfile.soft_reset()

def list_file(entry_list):
    print('\n[File Entries]')
    print('-' * 50)
    print('FileID Size  FileAttrib OffsetCopy0 OffsetCopy1')
    for entry in entry_list:
        print('0x%04X 0x%04X 0x%08X 0x%08X 0x%08X' % (
            entry.FileID,
            entry.FileSize,
            entry.FileAttrib,
            entry.OffsetCopy0,
            entry.OffsetCopy1)
        )

def read_file(entry_list, dir_name, file_type):
    error = []
    print('\n[Read File]')
    print('-' * 50)
    for entry in entry_list:
        file_id = entry.FileID
        sector_cnt = entry.FileSize if file_type == '.rpm' else 0
        if file_id == 0:
            continue
        if file_id == 1 and file_type == '.bin':
            continue
        print('0x%04X - ' % file_id, end='')
        start = time()
        if rdfl(file_id, sector_cnt):
            filename = '%s\\%X%s' % (dir_name, file_id, file_type)
            svbin(filename)
            end = time() - start
            print('%s %.2fs' % (dec_err_code(), end))
        else:
            error.append(file_id)
            errinfo('%s' % dec_err_code())
    if error:
        errinfo('Error(%d): '%len(error)+''.join(['0x%04X '% x for x in error]))

def write_file(entry_list, dir_name, file_type):
    error = []
    print('\n[Write File]')
    print('-' * 50)
    for entry in entry_list:
        file_id = entry.FileID
        sector_cnt = entry.FileSize if file_type == '.rpm' else 0
        if file_id == 0:
            continue
        if file_id == 1 and file_type == '.bin':
            continue
        filename = '%s\\%X%s' % (dir_name, file_id, file_type)
        if ldbin(filename):
            print('0x%04X - ' % file_id, end='')
            start = time()
            if not wrfl(file_id, sector_cnt):
                error.append(file_id)
                errinfo('%s' % dec_err_code())
            else:
                end = time() - start
                print('%s %.2fs' % (dec_err_code(), end))
        else:
            errinfo('0x%04X - No File: %s' % (file_id, filename))
    if error:
        errinfo('Error(%d): '%len(error)+''.join(['0x%04X '% x for x in error]))

def check_flash_file(entry_list):
    error = []
    print('\n[Check File]')
    print('-' * 50)
    for entry in entry_list:
        file_id = entry.FileID
        if file_id == 0 or file_id == 1:
            continue
        start = time()
        if rdfl(file_id, 1):
            end = time() - start
            print('0x%04X - SUCCESS %.2fs' % (file_id, end))
        else:
            errinfo('0x%04X - %s' % (file_id, dec_err_code()))
            error.append(file_id)
    if error:
        errinfo('Error(%d): '%len(error)+''.join(['0x%04X '% x for x in error]))

def check_res_file(entry_list):
    err_code = {
        '0': 'Good',
        '1': 'Invalid Header!',
        '2': 'Data Reading File!',
        '3': 'Servo Error Reading File!',
        '4': 'Compatible File!',
        '5': 'Checksums Different!',
        '6': 'Checksum Error!',
    }
    print('\n[Check File]')
    print('-' * 50)
    for entry in entry_list:
        file_id = entry.FileID
        if file_id == 0:
            continue
        print('Check File ID = 0x%04X ' % file_id, end='')
        start = time()
        if chkresf(file_id):
            end = time() - start
            print('%.2fs' % end)
            for hd in range(2):
                ecode = inp(ata_reg['low']['address']) if hd == 0 else \
                        inp(ata_reg['sc' ]['address'])
                ecode = str(ecode)
                des = err_code[ecode] if ecode in err_code else 'Unknown Error!'
                if ecode == '0':
                    print('    Hd%d: %s' % (hd, des))
                else:
                    errinfo('    Hd%d: %s' % (hd, des))
        else:
            errinfo('CmdFailed')
        cylx = inp(ata_reg['mid']['address'])+(inp(ata_reg['high']['address'])<<8)
        if cylx == 1:
            print('Check file block report is available. Please check read buffer!')

def backup_firmware():
    common.workdir = ''
    drv_tbl = get_drv_tbl()
    if common.workdir != '':
        path = common.get_startup_dir()
        common.backupdir = path + '\\Backup\\' + common.workdir
        if not os.path.exists(common.backupdir):
            os.makedirs(common.backupdir)
        common.workdir = path + '\\Work\\' + common.workdir
        if not os.path.exists(common.workdir):
            os.makedirs(common.workdir)
        filename = common.backupdir + '\\vscid.txt'
        if os.path.exists(filename):
            print('File Exist: ' + filename)
            return
        with open(filename, 'w', encoding='utf-8') as file:
            file.write(drv_tbl)
    flash_read(common.backupdir)
    flash_file_read(common.backupdir)
    res_read(common.backupdir)

def flash_read(dir_name=common.workdir):
    if read_flash(dir_name + '\\ROM.bin'):
        cprint('\nRead Flash - %s' % dec_err_code(), fore='light_green')
    else:
        cprint('\nRead Flash - %s' % dec_err_code(), fore='light_red')

def flash_write(dir_name=common.workdir):
    filename = dir_name + '\\ROM.bin'
    if not ldbin(filename, False):
        cprint('\nNo File - %s' % filename, fore='light_red')
        return
    if write_flash(filename):
        cprint('\nWrite Flash - %s' % dec_err_code(), fore='light_green')
    else:
        cprint('\nWrite Flash - %s' % dec_err_code(), fore='light_red')

def flash_file_read(dir_name=common.workdir):
    entry_list = get_flash_entry()
    list_file(entry_list)
    read_file(entry_list, dir_name, '.bin')

def flash_file_write(dir_name=common.workdir):
    entry_list = get_flash_entry()
    write_file(entry_list, dir_name, '.bin')

def flash_file_check():
    entry_list = get_flash_entry()
    check_flash_file(entry_list)

def res_read(dir_name=common.workdir):
    entry_list = get_file_entry()
    list_file(entry_list)
    read_file(entry_list, dir_name, '.rpm')

def res_write(dir_name=common.workdir):
    entry_list = get_file_entry()
    write_file(entry_list, dir_name, '.rpm')

def res_check():
    entry_list = get_file_entry()
    check_res_file(entry_list)

def res_read35(dir_name=common.workdir):
    print('Read Reserved File 0x35 - ', end='')
    start = time()
    if rdfl(0x35):
        filename = '%s\\%X%s' % (dir_name, 0x35, '.rpm')
        svbin(filename)
        end = time() - start
        print('%s %.2fs' % (dec_err_code(), end))
    else:
        errinfo('%s' % dec_err_code())

def load_ovl(dir_name):
    filename = dir_name + '\\11.rpm'
    if not os.path.exists(filename):
        errinfo('No File: %s' % filename)
        return False
    actions = [
        {'file': '10', 'fun': ld10, 'des': 'Load Cache Overlay(0x10)'},
        {'file': '5c', 'fun': ld5c, 'des': 'Load Servo Overlay(0x5C)'},
        {'file': '11', 'fun': ld11, 'des': 'Load Permanent Overlay(0x11)'},
        {'file': '13', 'fun': ld13, 'des': 'Load Cache Overlay(0x13)'},
        {'file': '80', 'fun': ld80, 'des': 'Load Overlay File(0x80)'},
        {'file': '1b', 'fun': ld1b, 'des': 'Load Transient Overlay(0x1B)'},
    ]
    result = False
    for i in range(len(actions)):
        file_id = actions[i]['file']
        if os.path.exists(dir_name + '\\' + file_id + '.rpm'):
            print(actions[i]['des'] + ' - ', end='')
            result = actions[i]['fun'](dir_name)
            if result:
                print('%s' % dec_err_code())
            else:
                get_err_code()
                errinfo('%s' % dec_err_code())
                break
    return result

def load_ovl_wrfl(dir_name=common.workdir):
    if not load_ovl(dir_name):
        return
    filename = dir_name + '\\1.rpm'
    if ldbin(filename):
        print('Write Reserved File 0x1 - ', end='')
        if not wrfl(1):
            errinfo('%s' % dec_err_code())
            return
        else:
            print('%s' % dec_err_code())
    else:
        errinfo('No File: %s' % filename)
        return
    print('Down Load Static File 0x1 - ', end='')
    if not download_stfl(1):
        errinfo('%s' % dec_err_code())
        return
    print('SUCCESS')
    entry_list = get_file_entry()
    write_file(entry_list, dir_name, '.rpm')

def init_fw():
    import wdfile
    drv_tbl = Drv_Data_Tbl()
    if not get_drv_tbl_data(drv_tbl):
        return False
    fmt_type = drv_tbl.FormatType
    wdfile.init_fw(1, fmt_type)

